import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import os

def set_plot_style():
    """设置全局绘图样式"""
    sns.set(style="whitegrid", font_scale=1.2)  # 使用 seaborn 全局样式



def plot_feature_distribution(df, output_dir):
    """绘制数值特征的分布直方图"""
    numeric_features = df.select_dtypes(include=[np.number]).columns
    fig, axes = plt.subplots(4, 4, figsize=(20, 15))
    axes = axes.flatten()

    for i, col in enumerate(numeric_features[:16]):  # 显示前16个数值特征
        sns.histplot(df[col], ax=axes[i], kde=True)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "feature_distributions.png"))
    plt.close()


def plot_correlation_matrix(df, output_dir):
    """绘制特征相关性热力图"""
    numeric_df = df.select_dtypes(include=[np.number])
    corr = numeric_df.corr()
    plt.figure(figsize=(12, 10))
    sns.heatmap(corr, annot=False, cmap='coolwarm', fmt=".2f", linewidths=.5)
    plt.title("Feature Correlation Matrix")
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "correlation_matrix.png"))
    plt.close()


def plot_true_vs_predicted(output_dir):
    """绘制测试集中真实值 vs 预测值对比图"""
    pred_df = pd.read_parquet("../data/output/predictions.parquet")
    plt.figure(figsize=(8, 6))
    sns.scatterplot(x="True_Value", y="Predicted_Value", data=pred_df, alpha=0.6)
    plt.plot([pred_df["True_Value"].min(), pred_df["True_Value"].max()],
             [pred_df["True_Value"].min(), pred_df["True_Value"].max()],
             'r--')  # 添加理想线
    plt.title("True vs Predicted Values (Test Set)")
    plt.xlabel("True Value")
    plt.ylabel("Predicted Value")
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "true_vs_predicted.png"))
    plt.close()


def plot_feature_importance(output_dir):
    """加载并显示特征重要性图"""
    importance_path = "../data/output/feature_importance.png"
    if os.path.exists(importance_path):
        plt.figure(figsize=(12, 8))
        img = plt.imread(importance_path)
        plt.imshow(img)
        plt.axis('off')
        plt.title("Top Feature Importances")
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, "feature_importance.png"))
        plt.close()
    else:
        print("⚠ 特征重要性图不存在，请先运行特征工程模块")


def plot_model_mse_curve(output_dir):
    """加载并显示训练过程中的 MSE 曲线"""
    mse_path = "../data/output/rf_training_history.png"
    if os.path.exists(mse_path):
        plt.figure(figsize=(10, 5))
        img = plt.imread(mse_path)
        plt.imshow(img)
        plt.axis('off')
        plt.title("Validation MSE Across Folds")
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, "model_mse_curve.png"))
        plt.close()
    else:
        print(" MSE 曲线图不存在，请先运行模型训练模块")

def plot_tensorflow_loss_curve(output_dir):
    """绘制 TensorFlow 模型的 loss 曲线"""
    loss_path = "../data/output/tensorflow_training_history.png"
    if os.path.exists(loss_path):
        plt.figure(figsize=(10, 5))
        img = plt.imread(loss_path)
        plt.imshow(img)
        plt.axis('off')
        plt.title("TensorFlow Training Loss Curve")
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, "tensorflow_loss_curve.png"))
        plt.close()
    else:
        print("⚠️ TensorFlow loss 曲线不存在，请先运行深度学习模型训练")




if __name__ == "__main__":
    set_plot_style()
    viz_dir = "../data/output/visualizations"
    os.makedirs(viz_dir, exist_ok=True)

    print(" 正在加载数据...")
    raw_df = pd.read_parquet("../data/preprocessed/cleaned_data.parquet")
    engineered_df = pd.read_parquet("../data/features/feature_engineered_data.parquet")

    print(" 正在绘制特征分布图...")
    plot_feature_distribution(raw_df, viz_dir)

    print(" 正在绘制特征相关性热力图...")
    plot_correlation_matrix(raw_df, viz_dir)

    print(" 正在绘制真实值 vs 预测值对比图...")
    plot_true_vs_predicted(viz_dir)

    print(" 正在复制特征重要性图...")
    plot_feature_importance(viz_dir)

    print(" 正在复制模型训练曲线...")
    plot_model_mse_curve(viz_dir)

    print(f" 所有可视化图表已保存至：{viz_dir}")
    plot_tensorflow_loss_curve(viz_dir)